from torchvision.datasets import CIFAR10, CIFAR100
from torchvision.transforms import transforms
from torch.utils.data import DataLoader, Dataset
import torchvision
from .aug import TransformTwice


def get_densecl_loader(args):
    if args.dataset == 'cifar10':
        trainset = torchvision.datasets.CIFAR10(root=args.data_path, train=True, download=True,
                                                transform=TransformTwice(args))
        trainloader = DataLoader(
            trainset, batch_size=args.batch_size, drop_last=True, shuffle=True, num_workers=args.num_workers)
        return trainloader

    elif args.dataset == 'cifar100':
        trainset = torchvision.datasets.CIFAR100(root=args.data_path, train=True, download=True,
                                                 transform=TransformTwice(args))

        trainloader = DataLoader(
            trainset, batch_size=args.batch_size, drop_last=True, shuffle=True, num_workers=args.num_workers)

        return trainloader
    else:
        raise NotImplementedError